{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3d4a0ea1-960c-489f-ae5b-589d69313784",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dotenv:  E:\\work\\code\\ai\\.env\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "from openai import OpenAI\n",
    "from dotenv import load_dotenv, find_dotenv\n",
    "\n",
    "def get_openai_key():\n",
    "    print(\"dotenv: \",find_dotenv())\n",
    "    _ = load_dotenv(find_dotenv())\n",
    "    # api_key = os.environ['OPENAI_API_KEY']\n",
    "    return OpenAI(\n",
    "        api_key = os.environ.get(\"ARK_API_KEY\"),\n",
    "    base_url = \"https://ark.cn-beijing.volces.com/api/v3\",)\n",
    "\n",
    "client = get_openai_key()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f1de1b1e-536b-47c8-83db-b8cadf7899fe",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "中国的首都是北京。\n",
      "\n",
      "北京有着悠久的历史，是中国历史文化名城，曾是多个朝代的都城，承载着丰富的历史文化遗产，如故宫、天坛、长城等。\n",
      "\n",
      "如今，北京是中国的政治中心，是党和国家最高领导机关所在地；也是文化中心，汇聚众多文化机构、高等院校、科研院所等；还是国际交往中心，众多国际会议、外交活动在此举行 ；同时也是科技创新中心，在科技研发、创新产业等方面发挥着引领作用。  \n"
     ]
    }
   ],
   "source": [
    "def get_completion(prompt, model=\"ep-20250103110319-28xth\",temperature = 0):\n",
    "    message = [\n",
    "        {\"role\": \"user\", \"content\": prompt},\n",
    "    ]\n",
    "    completion = client.chat.completions.create( \n",
    "    model = model,  # your model endpoint ID\n",
    "    messages =message,\n",
    "    temperature=temperature\n",
    "    )\n",
    "    return completion.choices[0].message.content\n",
    "\n",
    "response = get_completion(\"中国的首都是哪里？\")\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "b68097e0-3631-4885-8460-c1024cdc4f2f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "To reverse the letters in \"lollipop\", we start from the last letter and work our way to the first.\n",
      "\n",
      "The original word is \"lollipop\".\n",
      "\n",
      "When we reverse it, we get \"popillol\".\n"
     ]
    }
   ],
   "source": [
    "response = get_completion(\"Taken the letters in lollipop \\\n",
    " and reverse them\")\n",
    "\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "13b9fdd3-919d-4349-940b-0eeba7d968ce",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "p-o-p-i-l-l-o-l\n"
     ]
    }
   ],
   "source": [
    "response = get_completion(\"\"\"Take the letters in \\\n",
    "l-o-l-l-i-p-o-p and reverse them\"\"\")\n",
    "\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "bf478140-ae13-438b-9055-69fc95bcd9b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_completion_from_messages(message, model=\"ep-20250103110319-28xth\",temperature = 0, max_tokens=500):\n",
    "    # message = [\n",
    "    #     {\"role\": \"user\", \"content\": prompt},\n",
    "    # ]\n",
    "    completion = client.chat.completions.create( \n",
    "    model = model,  # your model endpoint ID\n",
    "    messages =message,\n",
    "    temperature=temperature,\n",
    "    max_tokens=max_tokens\n",
    "    )\n",
    "    return completion.choices[0].message.content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "13a95fcf-2327-47cf-94f5-ce160499f77c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# 快乐的小鲸鱼\n",
      "在那深深大海里，\n",
      "有只快乐小鲸鱼。\n",
      "尾巴摆呀摆不停，\n",
      "海水也跟着嬉戏。\n",
      "\n",
      "泡泡咕噜咕噜起，\n",
      "像是彩色小惊喜。\n",
      "跃出水面瞧天际，\n",
      "阳光洒身多惬意。\n",
      "\n",
      "它和海豚捉迷藏，\n",
      "笑声藏在波浪里。\n",
      "游过珊瑚的丛林，\n",
      "快乐永在它心里 。 \n"
     ]
    }
   ],
   "source": [
    "message = [\n",
    "        {\"role\": \"system\", \"content\": '你是一个助理， 并以Seuss 苏斯博士的风格做出回答。'},\n",
    "        {\"role\": \"user\", \"content\": '就快乐的小鲸鱼为主题给我写一首短诗'},\n",
    "    ]\n",
    "response = get_completion_from_messages(message, temperature=1)\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "f63981db-8231-48b3-a062-bf8246f515d3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "在广阔的海洋里，小鲸鱼每天欢快地穿梭于五彩珊瑚间，与小鱼嬉戏、同海豚共舞，洋溢着无尽快乐的它在湛蓝世界中开启一段又一段奇妙旅程 。  \n"
     ]
    }
   ],
   "source": [
    "message = [\n",
    "        {\"role\": \"system\", \"content\": '你的所有答复只能是一句话。'},\n",
    "        {\"role\": \"user\", \"content\": '写一个关于快乐的小鲸鱼的故事'},\n",
    "    ]\n",
    "response = get_completion_from_messages(message, temperature=1)\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "f1f2a9f0-d983-4b21-bdc2-1f5102e9cdca",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "在那深深海洋里，小鲸鱼欢快游弋，快乐泡泡冒不停，开启奇妙好之旅。 \n"
     ]
    }
   ],
   "source": [
    "message = [\n",
    "        {\"role\": \"system\", \"content\": '你是一个助理， 并以Seuss 苏斯博士的风格做出回答。只回一句话'},\n",
    "        {\"role\": \"user\", \"content\": '写一个关于快乐的小鲸鱼的故事'},\n",
    "    ]\n",
    "response = get_completion_from_messages(message, temperature=1)\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "06689036-cc55-4b0b-ab9b-c60c9d7fbcbf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_completion_and_token(message, model=\"ep-20250103110319-28xth\",temperature = 0, max_tokens=500):\n",
    "    # message = [\n",
    "    #     {\"role\": \"user\", \"content\": prompt},\n",
    "    # ]\n",
    "    response = client.chat.completions.create( \n",
    "    model = model,  # your model endpoint ID\n",
    "    messages =message,\n",
    "    temperature=temperature,\n",
    "    max_tokens=max_tokens\n",
    "    )\n",
    "    content =  response.choices[0].message.content\n",
    "    token_dict = {\n",
    "    'prompt_tokens':response.usage.prompt_tokens,\n",
    "    'completion_tokens':response.usage.completion_tokens,\n",
    "    'total_tokens':response.usage.total_tokens,\n",
    "    }\n",
    "    return content, token_dict\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "c84d4618-e5b5-4193-8871-bf45720726cf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# 快乐的小鲸鱼\n",
      "在那深深大海里，\n",
      "有只小鲸超欢喜。\n",
      "蓝色波涛当滑梯，\n",
      "游来游去乐无比。\n",
      "\n",
      "泡泡串串往上冒，\n",
      "像是珍珠在舞蹈。\n",
      "尾巴甩呀甩高高，\n",
      "快乐音符海中飘。\n",
      "\n",
      "和那鱼儿交朋友，\n",
      "讲讲笑话乐悠悠。\n",
      "阳光洒下金辉柔，\n",
      "小鲸快乐无尽头 。 \n"
     ]
    }
   ],
   "source": [
    "messages =  [  \n",
    "{'role':'system', \n",
    " 'content':'你是一个助理， 并以 Seuss 苏斯博士的风格作出回答。'},    \n",
    "{'role':'user', \n",
    " 'content':'就快乐的小鲸鱼为主题给我写一首短诗'},  \n",
    "] \n",
    "response, token_dict = get_completion_and_token(messages)\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "e29e900a-d0cf-4fc0-8a72-2f829f45350d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'prompt_tokens': 43, 'completion_tokens': 88, 'total_tokens': 131}\n"
     ]
    }
   ],
   "source": [
    "print(token_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38efa828-f092-4782-a4a5-7fa9d7c0a8d0",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.21"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
